ANDROID: KVM: arm64: Fix TLB invalidation when coalescing into a block

Wnen coalescing a table into a block, the break-before-make sequence
must invalidate the whole range of addresses translated by the entry in
order to avoid the possibility of a TLB conflict.

Fix the coalescing post-table walker so that the whole range of the old
table is invalidated, rather than just the first address, since a
refcount of 1 on the child page is not sufficient to ensure the absence
of any valid mappings.

Cc: Sebastian Ene <sebastianene@google.com>
Reported-by: Mostafa Saleh <smostafa@google.com>
Fixes: 9e7e5db52c3a ("ANDROID: KVM: arm64: Coalesce host stage2 entries on ownership reclaim")
Bug: 331232642
Signed-off-by: Will Deacon <willdeacon@google.com>
Change-Id: I4c94f552e4385599ad88b1be50b69ffbafa64a9b
diff --git a/arch/arm64/kvm/hyp/pgtable.c b/arch/arm64/kvm/hyp/pgtable.c
index 64387388..ea51f04 100644
--- a/arch/arm64/kvm/hyp/pgtable.c
+++ b/arch/arm64/kvm/hyp/pgtable.c
@@ -887,7 +887,9 @@
 	 * of the page table page.
 	 */
 	if (mm_ops->page_count(childp) == 1) {
-		stage2_put_pte(ptep, data->mmu, addr, level, mm_ops);
+		kvm_clear_pte(ptep);
+		kvm_call_hyp(__kvm_tlb_flush_vmid, data->mmu);
+		mm_ops->put_page(ptep);
 		mm_ops->put_page(childp);
 	}
 }